import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import linregress, f_oneway
import pandas as pd
from itertools import combinations

T = 2000
shock_interval = 120
shock_magnitude = 1.2
N_runs = 50
noise_level = 0.02


def simulate_erf(seed):
    np.random.seed(seed)
    beta = np.zeros(T)
    sigma = np.zeros(T)
    lam = np.zeros(T)

    beta[0] = 0.2
    sigma[0] = 0.2
    lam[0] = 0.1

    for t in range(T - 1):

        if t > 0 and t % shock_interval == 0:
            sigma[t] += shock_magnitude

        lam[t+1] = lam[t] + 0.35 * sigma[t] - 0.10 * lam[t]
        beta[t+1] = beta[t] - lam[t] * sigma[t] - 0.05 * beta[t]

        sigma[t+1] = np.tanh(
            sigma[t] + 0.25 * beta[t+1]
        ) + noise_level * np.random.randn()

    return sigma, lam


def simulate_fixed(seed):
    np.random.seed(seed)
    beta = np.zeros(T)
    sigma = np.zeros(T)
    lam = np.ones(T) * 0.1

    beta[0] = 0.2
    sigma[0] = 0.2

    for t in range(T - 1):

        if t > 0 and t % shock_interval == 0:
            sigma[t] += shock_magnitude

        beta[t+1] = beta[t] - lam[t] * sigma[t] - 0.03 * beta[t]

        sigma[t+1] = np.tanh(
            sigma[t] + 0.25 * beta[t+1]
        ) + noise_level * np.random.randn()

    return sigma, lam


def simulate_exogenous(seed):
    np.random.seed(seed)
    beta = np.zeros(T)
    sigma = np.zeros(T)
    lam = np.zeros(T)

    beta[0] = 0.2
    sigma[0] = 0.2
    lam[0] = 0.1

    for t in range(T - 1):

        if t > 0 and t % shock_interval == 0:
            sigma[t] += shock_magnitude

        lam[t+1] = lam[t] + 0.4 * np.random.randn() - 0.08 * lam[t]
        beta[t+1] = beta[t] - lam[t] * sigma[t] - 0.05 * beta[t]

        sigma[t+1] = np.tanh(
            sigma[t] + 0.25 * beta[t+1]
        ) + noise_level * np.random.randn()

    return sigma, lam


def recovery_after_shock(series, window=50, tolerance=0.05):
    horizons = []
    for t in range(shock_interval, T, shock_interval):
        pre_mean = np.mean(series[t-window:t])
        for k in range(t, T):
            if abs(series[k] - pre_mean) < tolerance:
                horizons.append(k - t)
                break
    return np.mean(horizons)


def coupling_correlation(sigma, lam):
    _, _, r, _, _ = linregress(np.abs(sigma), lam)
    return r


def tail_variance(series):
    return np.var(series[-500:])


records = []

for run in range(N_runs):

    sigma_e, lam_e = simulate_erf(run)
    sigma_f, lam_f = simulate_fixed(run)
    sigma_x, lam_x = simulate_exogenous(run)

    for name, sigma, lam in [
        ("ERF", sigma_e, lam_e),
        ("Fixed", sigma_f, lam_f),
        ("Exogenous", sigma_x, lam_x)
    ]:
        records.append({
            "Run": run,
            "Configuration": name,
            "Recovery": recovery_after_shock(sigma),
            "CouplingCorrelation": coupling_correlation(sigma, lam),
            "TailVariance": tail_variance(sigma)
        })

df = pd.DataFrame(records)
df.to_csv("MonteCarlo_Raw_Noise.csv", index=False)

summary = df.groupby("Configuration").agg(
    Recovery_mean=("Recovery", "mean"),
    Recovery_std=("Recovery", "std"),
    Corr_mean=("CouplingCorrelation", "mean"),
    Corr_std=("CouplingCorrelation", "std"),
    Var_mean=("TailVariance", "mean"),
    Var_std=("TailVariance", "std")
).reset_index()

summary.to_csv("MonteCarlo_Summary_Noise.csv", index=False)


# ===============================
# Statistical Tests
# ===============================

def compute_anova(metric):
    groups = [df[df["Configuration"] == cfg][metric] for cfg in ["ERF","Fixed","Exogenous"]]
    F, p = f_oneway(*groups)
    return F, p


def cohens_d(x, y):
    nx = len(x)
    ny = len(y)
    pooled_std = np.sqrt(((nx-1)*np.var(x, ddof=1) + (ny-1)*np.var(y, ddof=1)) / (nx+ny-2))
    return (np.mean(x) - np.mean(y)) / pooled_std


anova_results = []
for metric in ["Recovery", "TailVariance"]:
    F, p = compute_anova(metric)
    anova_results.append({"Metric": metric, "F_stat": F, "p_value": p})

anova_df = pd.DataFrame(anova_results)
anova_df.to_csv("MonteCarlo_ANOVA.csv", index=False)


effect_records = []
for metric in ["Recovery", "TailVariance"]:
    for a, b in combinations(["ERF","Fixed","Exogenous"], 2):
        x = df[df["Configuration"] == a][metric]
        y = df[df["Configuration"] == b][metric]
        d = cohens_d(x, y)
        effect_records.append({
            "Metric": metric,
            "Comparison": f"{a} vs {b}",
            "Cohens_d": d
        })

effect_df = pd.DataFrame(effect_records)
effect_df.to_csv("MonteCarlo_EffectSizes.csv", index=False)


plt.figure(figsize=(8,4))
for name in ["ERF","Fixed","Exogenous"]:
    subset = df[df["Configuration"]==name]
    plt.hist(subset["Recovery"], bins=20, alpha=0.5, label=name)
plt.legend()
plt.savefig("MC_Recovery_Distribution.png")
plt.close()

plt.figure(figsize=(8,4))
for name in ["ERF","Fixed","Exogenous"]:
    subset = df[df["Configuration"]==name]
    plt.hist(subset["TailVariance"], bins=20, alpha=0.5, label=name)
plt.legend()
plt.savefig("MC_Variance_Distribution.png")
plt.close()

print(summary)
print(anova_df)
print(effect_df)